#!/usr/bin/env python2.7

# lidar_inspector.py

import os, sys
#import re
from datetime import datetime, timedelta
#from dateutil.parser import parse
import numpy as np
import matplotlib
#matplotlib.use('Agg')
import matplotlib.pyplot as plt
from  matplotlib.dates import num2date, date2num
from matplotlib import transforms
from mpl_toolkits.basemap.cm import *
#from matplotlib import mlab
#from netCDF4 import Dataset
from metlib.lidar import *
from metlib.color import *
from metlib.shell import *
from metlib.datetime import *
from optparse import OptionParser

class LidarInspector(object):
    def __init__(self, filelist, rc_file=None, max_height=5000.0, aver_min=0, datafield='data/energy', channel=0):
        """
        datafield: 'data', 'raw', 'data/energy'
        """
        self.filelist = filelist
        self.now_file_i = 0
        self.datafield = datafield
        print self.datafield
        self.channel = channel
        self.max_height = max_height
        self.aver_min = aver_min
        print self.filelist
        self.cmds = []
        self._init_cmaps()
        self.load_lidar_data(self.now_file_i)
        self._init_plot()
        self.when_new_data_loaded()
        
        self.set_keymap()
        self.set_reg()
        if rc_file is not None:
            f = open(rc_file)
            for line in f:
                line = line.rstrip()
                exec(line)

    def _init_plot(self):
        self.p = LidarPlot(self.dts_ext, self.height_ext, self.data[:,self.channel,:], vmin=-3.0, vmax=1.0, title=self.get_status_str(), xlabel=self.dts[0].strftime('%Y-%m-%d'))
        self.p.hax.set_ylim(-3.0, 1.0)
        self.p.vax.set_xlim(-3.0, 1.0)
        self.tb = plt.get_current_fig_manager().toolbar
        self.hplot = self.p.hax.plot(self.dtsnum, self.hzeros, color='black', visible=True)[0]
        self.vplot = self.p.vax.plot(self.vzeros, self.height, color='black', visible=True)[0]
        self.plot_hvlines()
        self.pointer_i = 0
        self.reg_event()

    def reg_event(self):
        self.p.fig.canvas.mpl_connect('button_press_event', self.on_click)
        self.p.fig.canvas.mpl_connect('key_press_event', self.on_key)


    def plot_hvlines(self):
        self.hline = self.p.ax.axhline(0, color='gray', visible=True)
        self.vline = self.p.ax.axvline(self.dtsnum[0], color='gray', visible=True)

    def _init_cmaps(self):
        self.cmaps = {'1':matplotlib.cm.jet,
                    '2':matplotlib.cm.Accent,
                    '3':matplotlib.cm.Paired,
                    '4':matplotlib.cm.gist_ncar,
                    '5':cm_pop_poster_r,
                    '6':matplotlib.cm.spectral,
                    '7':cm_GateRdBu1_r,
                    '&':cm_GateBuRd,
                    '8':cm_GateBuRd2,
                    '9':matplotlib.cm.RdYlBu_r,
                    '0':s3pcpn_r,
                    ')':s3pcpn,
                    }

    def get_status_str(self):
        return '%s %s' % (self.status1, self.status2)
    
    def load_file(self, file_no):
        tmp_file_i = self.now_file_i
        if file_no in ('n', 'p'):
            if file_no == 'n':
                if tmp_file_i == len(self.filelist) - 1:
                    return
                tmp_file_i += 1
            elif file_no == 'p':
                if tmp_file_i == 0:
                    return
                tmp_file_i -= 1
        else:
            if file_no < 0 and file_no >= -len(self.filelist):
                tmp_file_i = len(self.filelist) + file_no
            elif file_no < len(self.filelist) and file_no >= 0:
                tmp_file_i = file_no
            else:
                return
        print self.now_file_i
        self.now_file_i = tmp_file_i
        self.load_lidar_data(self.now_file_i)
        self.when_new_data_loaded()
        plt.draw()
    
    def load_lidar_data(self, i):
        print "Load_lidar_data: %s" % self.filelist[i]
        try:
            to_load = self.filelist[i]
            if isinstance(to_load, (str, unicode)):
                self.d = LidarDataset(to_load)
            else:
                self.d = to_load
            max_i = height_to_index(self.max_height, self.d, 75.0)
            self.d.resize_bin(0, max_i)
            if self.aver_min != 0:
                self.d.time_average(self.aver_min)
            self.update_lidar_info()
        except Exception as e:
            print e

    def update_lidar_info(self):
        self.dts = self.d['datetime']
        self.tdelta = self.dts[-1] - self.dts[-2]
        self.dts_off = self.dts + self.tdelta / 2
        self.dts_ext = np.hstack((self.dts, self.dts[-1] + self.tdelta))
        self.dtsnum = date2num(self.dts)
        self.tdnum = self.dtsnum[-1] - self.dtsnum[-2]
        self.dtsnum_ext = date2num(self.dts_ext)
        self.dtsnum_off = date2num(self.dts_off)
        self.height = self.d['distance'] * np.sin(np.deg2rad(75))
        self.height_ext = np.hstack((self.height, self.height[-1] + (self.height[-1] - self.height[-2])))
        if self.datafield == 'data/energy':
            self.raw = (self.d['data'] / self.d['energy'][..., np.newaxis])
        else:
            self.raw = self.d[self.datafield]
        if self.d['lidarname'] == "Yuenlong Lidar":
            start_i = 6
        elif self.d['lidarname'] == "ShaTauKok Lidar":
            if self.d['bin_size'] < 10.0:
                start_i = 24
            else:
                start_i = 12
        elif self.d['lidarname'] == "Pump House Lidar":
            start_i = 250
        else:
            start_i = 0
        self.raw[...,:start_i] = self.raw[...,start_i:start_i+1]
        self.data = np.log10(self.raw)
        self.data_bak = self.data.copy()
        self.status1 = 'log10(RCS)'
        self.status2 = ''
        self.hzeros = np.zeros_like(self.dtsnum)
        self.vzeros = np.zeros_like(self.height)

    def when_new_data_loaded(self):
        # cmds
        for cmd in self.cmds:
            self.on_key(cmd, reg_cmd=False)
        # redraw pcolormesh
        self.p.update_data(self.dts_ext, self.height_ext, self.data[:,self.channel,:])
        self.plot_hvlines()
        # xlim
        self.p.ax.set_xlim(self.dtsnum[0], self.dtsnum_ext[-1])
        self.p.hax.set_xlim(self.dtsnum[0], self.dtsnum_ext[-1])
        self.p.suptitle.set_text(self.get_status_str())
        self.p.xlabel.set_text(self.dts[0].strftime('%Y-%m-%d'))

    def xy2ij(self, x, y):
        ix = np.argmin(np.abs(x-self.tdnum/2 - self.dtsnum))
        jy = np.argmin(np.abs(y - self.height))
        return ix, jy

    def on_click(self, event):
        if self.tb.mode != '':
            return
        if event.inaxes != self.p.ax:
            return
#        print event.__dict__
#        print num2date(event.xdata), event.ydata
        i, j = self.xy2ij(event.xdata, event.ydata)
#        print i, j, self.dtsnum[i]
        self.pointer_i = i
        self.hline.set_data([0,1], [self.height[j], self.height[j]])
        self.vline.set_data(np.tile(self.dtsnum[i]+self.tdnum/2, 2), [0,1])
        self.hplot.set_data(self.dtsnum, self.data[:,self.channel,j])
        self.vplot.set_data(self.data[i,self.channel,:], self.height)
        self.on_click_user(event)
        plt.draw()
    
    def on_click_user(self, event):
        pass

    def on_release(self, event):
        if event.inaxex != self.p.ax: return
        if self.tb.mode != '':
            self.p.update()
    
    def on_key(self, event, reg_cmd=True):
#        print event.key
        if event.key in self.keymap:
            cmd_tuples = self.keymap[event.key]
#            print cmd_tuples
            for cmd_bundle in cmd_tuples:
                cmd_bundle[0](*cmd_bundle[1:])
            if reg_cmd:
                if event.key in self.keys_to_reg:
                    self.cmds.append(event)
                elif event.key in self.keys_to_reset_reg:
                    self.cmds = []
                elif event.key in self.keys_to_reset_and_reg:
                    self.cmds = [event]

            
    def set_image(self, arr=None, status1=None, status2=None):
        if arr is None:
            arr = self.data[:,self.channel,:]
        if status1 is not None:
            self.status1 = status1
        if status2 is not None:
            self.status2 = status2
        self.p.image.set_array(np.ma.masked_invalid(arr.transpose().flatten()))
        self.p.suptitle.set_text(self.get_status_str())
        plt.draw()

    def set_data_mode(self, mode, status1=None, status2=None, *args):
        self.data_bak = self.data.copy()
        if callable(mode):
            self.data = mode(self, *args)
        if mode == 'RCS':
            self.data = self.raw.copy()
            self.status1 = 'RCS'
            self.status2 = ''
        elif mode == 'log10(RCS)':
            self.data = np.log10(self.raw)
            self.status1 = 'log10(RCS)'
            self.status2 = ''

        if status1 is not None:
            self.status1 = status1
        if status2 is not None:
            if status2.startswith('+'):
                self.status2 += status2
            else:
                self.status2 = status2

    def set_cmap(self, cmap):
        if isinstance(cmap, (str, unicode)):
            cmap = self.cmaps[cmap]
        self.p.image.set_cmap(cmap)
        plt.draw()
    
    def set_lim(self, ax_from, ax_to):
        ilim = self.p.image.get_clim()
        hlim = self.p.hax.get_ylim()
        vlim = self.p.vax.get_xlim()
        lim_from = {'ax':ilim, 'hax':hlim, 'vax':vlim}[ax_from]
        if ax_to == 'ax':
            self.p.image.set_clim(lim_from)
        elif ax_to == 'hax':
            self.p.hax.set_ylim(lim_from)
        elif ax_to == 'vax':
            self.p.vax.set_xlim(lim_from)
        plt.draw()

    def set_vrange(self, min_change, max_change):
        now_min, now_max = self.p.image.get_clim()
        delta = (now_max - now_min) * 0.02
        delta_large = delta * 5
        change_table = {'+':delta, '-':-delta, '++':delta_large, '--':-delta_large, '':0}
        min_inc = change_table[min_change]
        max_inc = change_table[max_change]
        now_min += min_inc
        now_max += max_inc
        if now_min >= now_max:
            now_min = now_max - delta
        self.p.image.set_clim(now_min, now_max)
        plt.draw()

    def set_reg(self, keys=None, update=True, reg_type='keys_to_reg'):
        if keys is not None:
            if update:
                self.__dict__[reg_type].update(keys)
            else:
                self.__dict__[reg_type] = set(keys)
        else:
            self.keys_to_reg = set([])
            self.keys_to_reset_reg = set(['f1'])
            self.keys_to_reset_and_reg = set(['f2'])
        

    def add_data_keymap(self, key, func, status1, status2, *args):
        # TODO: test this function
        cmd = [self.set_data_mode, func, status1, status2]
        cmd.extend(args)
        self.keymap[key] = (tuple(cmd), (self.set_image,))

    def set_keymap(self, keymap=None, update=True):
        if keymap is not None:
            if update:
                self.keymap.update(keymap)
            else:
                self.keymap = keymap
        else:
            self.keymap = {
                    # file
                    'pageup':((self.load_file, 'p') , (self.set_image,)) ,
                    'pagedown':((self.load_file, 'n') ,(self.set_image,)) ,
                    # ax value lim
                    'H':((self.set_lim, 'ax', 'hax'),),
                    'J':((self.set_lim, 'hax', 'ax'),),
                    'C':((self.set_lim, 'ax', 'vax'),),
                    'V':((self.set_lim, 'vax', 'ax'),),
                    # value range
                    'q':((self.set_vrange, '+', ''),),
                    'a':((self.set_vrange, '-', ''),),
                    'e':((self.set_vrange, '', '+'),),
                    'd':((self.set_vrange, '', '-'),),
                    'Q':((self.set_vrange, '++', ''),),
                    'A':((self.set_vrange, '--', ''),),
                    'E':((self.set_vrange, '', '++'),),
                    'D':((self.set_vrange, '', '--'),),
                    '[':((self.set_vrange, '-', '-'),),
                    ']':((self.set_vrange, '+', '+'),),
                    '{':((self.set_vrange, '--', '--'),),
                    '}':((self.set_vrange, '++', '++'),),
                    # mode
                    'f1':((self.set_data_mode, 'log10(RCS)', 'log10(RCS)', '') ,(self.set_image,)) ,
                    'f2':((self.set_data_mode, 'RCS', 'RCS', '') ,(self.set_image,)) ,
#                    'f4':((self.set_data_mode, do_log10, None, '+log10'), (self.set_image,)),
                    # algorithm
                }
            for c in '1234567890&)':
                self.keymap[c] = ((self.set_cmap, c), )

    def set_reg(self, keys=None, update=True, reg_type='keys_to_reg'):
        if keys is not None:
            if update:
                self.__dict__[reg_type].update(keys)
            else:
                self.__dict__[reg_type] = set(keys)
        else:
            self.keys_to_reg = set([])
            self.keys_to_reset_reg = set(['f1'])
            self.keys_to_reset_and_reg = set(['f2'])
        

if __name__ == '__main__':
    parser = OptionParser()
    parser.add_option('-v', '--average', dest='aver_min', type='int',
            default=0, metavar='Average Minutes(default 0 means no average)')
    (options, args) = parser.parse_args()
    print options.aver_min
    fnames = args
    local_rc = './dodo_lidar_browser.rc'
    user_rc = os.path.expanduser('~/.dodo_lidar_browser.rc')
    if os.path.exists(local_rc):
        rc_file = local_rc
    elif os.path.exists(user_rc):
        rc_file = user_rc
    else:
        rc_file = None

    li = LidarInspector(filelist=fnames, rc_file=rc_file, max_height=5000.0, aver_min=options.aver_min)
    plt.show()
